# detection_server.py
import socket
import threading
import tkinter as tk
from tkinter import scrolledtext
import json
import os
from datetime import datetime

class DetectionResultServer:
    def __init__(self, host="127.0.0.1", port=5001): # 使用 5001 端口以避免与 Flask 冲突
        self.clients = []
        self.host = host
        self.port = port
        self.lock = threading.Lock() # 保护 clients 列表的线程安全
        self.results_dir = "received_results"
        os.makedirs(self.results_dir, exist_ok=True) # 确保保存目录存在

        # --- 新增：服务器暂停状态标志 ---
        self.is_paused = False
        # --- 新增结束 ---

        # --- 新增：用于控制空数据日志频率 ---
        self.empty_data_counter = 0
        self.log_empty_every_n = 10 # 每收到 N 次空数据才记录一次
        # --- 新增结束 ---

        # 图形界面
        self.window = tk.Tk()
        self.window.title("YOLO 检测结果接收服务器")

        self.text_area = scrolledtext.ScrolledText(self.window, wrap=tk.WORD, height=25, width=80, state="disabled")
        self.text_area.pack(padx=10, pady=10)

        # --- 修改：创建按钮框架以更好地布局按钮 ---
        button_frame = tk.Frame(self.window)
        button_frame.pack(pady=5)

        self.start_button = tk.Button(button_frame, text="启动服务", command=self.start_server)
        self.start_button.pack(side=tk.LEFT, padx=5)
        
        # --- 新增：创建暂停按钮 ---
        self.pause_button = tk.Button(button_frame, text="暂停服务", command=self.toggle_pause, state="disabled")
        self.pause_button.pack(side=tk.LEFT, padx=5)
        # --- 新增结束 ---
        
        self.status_label = tk.Label(self.window, text="状态: 未启动")
        self.status_label.pack()

    # --- 新增：切换暂停状态的方法 ---
    def toggle_pause(self):
        """切换服务器的暂停/恢复状态"""
        with self.lock: # 使用锁保护 is_paused 状态的修改
            self.is_paused = not self.is_paused
        
        if self.is_paused:
            self.log_message("[服务器] 服务已暂停。")
            self.status_label.config(text=f"状态: 已暂停 (监听 {self.host}:{self.port})")
            self.pause_button.config(text="恢复服务")
        else:
            self.log_message("[服务器] 服务已恢复。")
            self.status_label.config(text=f"状态: 运行中 (监听 {self.host}:{self.port})")
            self.pause_button.config(text="暂停服务")
    # --- 新增结束 ---

    def start_server(self):
        """启动服务器"""
        try:
            self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) # 允许地址重用
            self.server.bind((self.host, self.port))
            self.server.listen(5)
            self.log_message(f"[服务器] 已启动，监听 {self.host}:{self.port}...")
            self.status_label.config(text=f"状态: 运行中 (监听 {self.host}:{self.port})")
            self.start_button.config(state="disabled", text="服务运行中") # 启动后禁用按钮

            # --- 新增：启用暂停按钮并初始化暂停状态 ---
            self.pause_button.config(state="normal", text="暂停服务")
            self.is_paused = False # 启动时服务器不暂停
            # --- 新增结束 ---

            # 在新线程中接受客户端连接
            thread = threading.Thread(target=self.accept_clients, daemon=True)
            thread.start()
        except Exception as e:
            self.log_message(f"[错误] 启动服务器失败: {e}")
            self.status_label.config(text="状态: 启动失败")

    def accept_clients(self):
        """接受客户端连接"""
        while True:
            # --- 新增：检查暂停状态 ---
            with self.lock:
                if self.is_paused:
                    # 如果暂停，则短暂休眠并继续循环（不接受新连接）
                    # 使用 wait() 和 Condition 更高效，但 sleep 足够简单
                    self.server._sock.close() # 关闭监听 socket
                    break # 退出循环，停止接受新连接
            # --- 新增结束 ---
            
            # 注意：上面的实现有个问题，就是一旦暂停就无法恢复 accept。
            # 更好的方式是在循环内部检查，并使用条件变量或者事件来控制。
            # 但是 socket.accept() 是阻塞的，不容易中断。
            # 一个更简单的折衷方案是：暂停时仍然 accept，但在 handle_client 中不处理数据。
            # 但我们按原意实现：暂停 accept。需要在 toggle_pause 中重新 listen。
            # 让我们修正这个逻辑。
            
            # 修正后的逻辑：在 toggle_pause 中完全 stop 并 restart server socket 更复杂。
            # 我们采用检查 is_paused 的方式，但 accept 本身无法轻易中断。
            # 一个可行的办法是设置一个超时，定期检查 is_paused。
            # 但最简单的是在 handle_client 中处理 is_paused。
            # 这里我们保持 accept 运行，主要在 handle_client 中控制处理。
            # 因此，暂停主要影响 handle_client 和新功能的开启，不影响已建立的连接。
            # 为了简单起见，我们主要在 handle_client 中实现暂停效果。
            # 这里的 accept 暂停逻辑需要重新考虑。
            # 让我们简化：暂停按钮只影响数据处理（handle_client），不停止 accept。
            # 这样实现最简单且符合大部分需求。
            
            # 恢复原始逻辑，并在 handle_client 中检查 is_paused
            # --- 删除上面有问题的代码块 ---
            pass
            # --- 删除结束 ---

            try:
                # 设置一个较短的超时，以便能定期检查 is_paused
                # 注意：timeout 是浮点数，单位秒
                self.server.settimeout(1.0) # 1秒超时
                client, addr = self.server.accept()
                self.server.settimeout(None) # 重置为阻塞模式
                
                with self.lock:
                    self.clients.append(client)
                self.log_message(f"[连接] 客户端 {addr} 已连接。")
                # 为每个客户端创建一个处理线程
                threading.Thread(target=self.handle_client, args=(client, addr), daemon=True).start()
            except socket.timeout:
                # 超时，检查是否需要暂停 accept 线程
                # 这里 accept 线程继续循环，主要暂停效果在 handle_client 中体现
                continue
            except OSError as e:
                # 如果服务器 socket 被关闭（例如在 toggle_pause 中）
                if "Bad file descriptor" in str(e) or "Invalid argument" in str(e):
                     self.log_message("[服务器] 接受连接线程已停止。")
                     break
                else:
                     self.log_message(f"[错误] 接受客户端连接时出错 (OSError): {e}")
                     break
            except Exception as e:
                # 如果服务器 socket 被关闭，accept 会抛出异常
                if hasattr(self, 'server'):
                    self.log_message(f"[错误] 接受客户端连接时出错: {e}")
                break

    def handle_client(self, client, addr):
        """处理单个客户端的消息"""
        while True:
            # --- 新增：在处理循环开始时检查暂停状态 ---
            # 注意：在暂停状态下，线程仍然存活，但不会处理新数据
            paused = False
            with self.lock:
                paused = self.is_paused
            if paused:
                # 可以选择在这里短暂休眠，避免空转
                # self.log_message(f"[调试] 客户端 {addr} 处理线程已暂停...") # 调试用，会很频繁
                threading.Event().wait(0.1) # 等待 100ms
                continue # 继续下一次循环检查
            # --- 新增结束 ---

            try:
                # 设置接收超时，以便能定期检查 is_paused
                client.settimeout(1.0) # 1秒超时
                data_bytes = client.recv(4096)
                client.settimeout(None) # 重置为阻塞模式

                if not data_bytes:
                    break # 客户端断开连接

                # 解码字节数据为字符串
                data_str = data_bytes.decode('utf-8')
                # 可选：记录接收到原始数据（非常频繁，通常注释掉）
                # self.log_message(f"[接收] 来自 {addr} 的原始数据: {data_str[:100]}...") # 只记录前100字符

                # 尝试解析 JSON
                try:
                    data_dict = json.loads(data_str)
                    
                    # --- 修改：明确检查 xywh 字段 ---
                    # 安全地获取 xywh 和 class 字段
                    xywh_data = data_dict.get("xywh", "") # 默认为空字符串
                    class_data = data_dict.get("class", "") # 默认为空字符串

                    # --- 修改：明确检查 xywh 是否为非空字符串 ---
                    if xywh_data: # 检查 xywh 是否为 True (非空字符串)
                        # --- 处理非空 xywh 数据 ---
                        # (处理逻辑保持不变)
                        try:
                            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
                            filename = f"result_{timestamp}.json"
                            filepath = os.path.join(self.results_dir, filename)
                            
                            with open(filepath, 'w', encoding='utf-8') as f:
                                json.dump(data_dict, f, indent=4, ensure_ascii=False)
                            
                            self.log_message(f"[保存] 非空结果已保存: {filepath} (class='{class_data}', xywh_len={len(xywh_data)})")
                            ack_message = json.dumps({"status": "success", "message": "Non-empty result received and saved"}, ensure_ascii=False)
                            
                        except Exception as save_err:
                            error_msg = f"[错误] 保存文件失败 ({filename}): {save_err}"
                            self.log_message(error_msg)
                            ack_message = json.dumps({"status": "error", "message": f"Save failed: {str(save_err)}"}, ensure_ascii=False)
                        
                        client.send(ack_message.encode('utf-8'))
                        # --- 处理非空 xywh 数据结束 ---
                    
                    else:
                        # --- 处理 xywh 为空的数据 ---
                        self.empty_data_counter += 1
                        
                        # --- 修改：降低空数据日志频率 ---
                        if self.empty_data_counter % self.log_empty_every_n == 0:
                             self.log_message(f"[接收] 来自 {addr} 的空数据 (累计 {self.empty_data_counter} 次)")
                        # --- 修改结束 ---
                        
                        ack_message = json.dumps({"status": "success", "message": "Empty result received"}, ensure_ascii=False)
                        client.send(ack_message.encode('utf-8'))
                        # --- 处理 xywh 为空的数据结束 ---

                except json.JSONDecodeError as json_err:
                    warning_msg = f"[警告] 来自 {addr} 的数据不是有效的 JSON: {data_str[:100]}..."
                    self.log_message(warning_msg)
                    error_message = json.dumps({"status": "error", "message": "Invalid JSON"}, ensure_ascii=False)
                    try:
                        client.send(error_message.encode('utf-8'))
                    except Exception as send_err:
                        self.log_message(f"[错误] 向 {addr} 发送 JSON 错误回复失败: {send_err}")
                # --- 解析 JSON 结束 ---

            except socket.timeout:
                # 接收超时，继续下一次循环检查 is_paused
                continue
            except ConnectionResetError:
                self.log_message(f"[断开] 客户端 {addr} 强制断开连接。")
                break
            except Exception as e:
                error_msg = f"[严重错误] 处理客户端 {addr} 时出错: {e}"
                self.log_message(error_msg)
                import traceback
                self.log_message(traceback.format_exc())
                break

        # 客户端断开连接后的清理工作
        try:
            client.close()
        except Exception as close_err:
            self.log_message(f"[错误] 关闭客户端 {addr} 连接时出错: {close_err}")
        with self.lock:
            if client in self.clients:
                self.clients.remove(client)
        self.log_message(f"[断开] 客户端 {addr} 连接已关闭。")


    def log_message(self, message):
        """线程安全地更新 GUI 日志"""
        self.window.after(0, self._log_message_safe, message)

    def _log_message_safe(self, message):
        """在主线程中实际更新 GUI"""
        self.text_area.config(state="normal")
        self.text_area.insert(tk.END, f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {message}\n")
        self.text_area.config(state="disabled")
        self.text_area.see(tk.END)

    def run(self):
        """运行 GUI 主循环"""
        self.window.protocol("WM_DELETE_WINDOW", self.on_closing)
        self.window.mainloop()
        
    def on_closing(self):
        """处理窗口关闭事件"""
        self.log_message("[服务器] 正在关闭...")
        self.status_label.config(text="状态: 正在关闭...")
        
        # --- 修改：确保在关闭时也能正确清理 ---
        with self.lock:
             self.is_paused = False # 确保 accept 线程可以退出
        # --- 修改结束 ---
        
        # 关闭所有客户端连接
        with self.lock:
            for client in self.clients[:]:
                try:
                    client.close()
                except Exception as close_err:
                    self.log_message(f"[错误] 关闭客户端连接时出错: {close_err}")
            self.clients.clear()
        
        # 关闭服务器 socket
        try:
            if hasattr(self, 'server'):
                self.server.close()
        except Exception as server_close_err:
            self.log_message(f"[错误] 关闭服务器 socket 时出错: {server_close_err}")
        
        self.log_message("[服务器] 已关闭。")
        self.status_label.config(text="状态: 已关闭")
        # --- 修改：重置按钮状态 ---
        self.start_button.config(state="normal", text="启动服务")
        self.pause_button.config(state="disabled", text="暂停服务")
        # --- 修改结束 ---
        self.window.destroy()

if __name__ == "__main__":
    server = DetectionResultServer()
    server.run()